#!/bin/bash
set -e

# 配置参数
KEY="secret_key"  # 加密密钥（与loader.c中的密钥一致）
OUTPUT_MAIN_A="$1"  # 主程序名称
OUTPUT_LOADER="$2"  # 最终生成的加载器名称
ENCRYPTED_BIN="main_a.enc"    # 加密后的二进制文件
IV_FILE="iv.bin"              # 初始化向量文件
LOADER_C="new_loader.c"
MAGIC="ATTACHEND"             # 附加数据的魔数

# 检查参数
if [ -z "$OUTPUT_MAIN_A" ] || [ -z "$OUTPUT_LOADER" ] || [ -z "$3" ] || [ -z "$4" ]; then
  echo "使用方法: $0 <主程序名称> <输出加载器名称> <编译器路径> <架构(x86|aarch64)>"
  exit 1
fi

echo "[*] Step 1: 使用 AES-256-CBC 加密 $OUTPUT_MAIN_A"
echo "[*] 加密密钥: '$KEY'"

# 生成随机初始化向量 (IV)
openssl rand -out "$IV_FILE" 16

# 使用 OpenSSL 加密文件
openssl enc -aes-256-cbc \
    -in "$OUTPUT_MAIN_A" \
    -out "$ENCRYPTED_BIN" \
    -K $(printf "$KEY" | sha256sum | cut -d' ' -f1) \
    -iv $(xxd -p "$IV_FILE" | tr -d '\n') \
    -nosalt

echo "[*] Step 2: 合并 IV 和加密数据"
# 合并 IV 和加密数据
cat "$IV_FILE" "$ENCRYPTED_BIN" > combined.bin

# 获取加密后的文件大小
PAYLOAD_LEN=$(stat -c%s combined.bin)

echo "[*] Step 3: 编译基础 loader"
CC="$3"
if [ "$4" == "x86" ]; then
  INCS=$RSDP_DIR2/inc/third/x86/standard/libopenssl/
  LIBS=$RSDP_DIR2/x86/standard/libthird
elif [ "$4" == "aarch64" ]; then
  INCS=$RSDP_DIR2/inc/third/arm/aarch64/libopenssl/
  LIBS=$RSDP_DIR2/arm/aarch64/libthird/
else
  echo "错误: 不支持的架构 '$4'，仅支持 x86 或 aarch64"
  exit 1
fi

# 编译基础 loader
$CC -s -O2 -o "$OUTPUT_LOADER" "$LOADER_C" -L "$LIBS" -lcrypto -ldl -I "$INCS"

# 检查编译是否成功
if [ ! -f "$OUTPUT_LOADER" ]; then
  echo "错误: 编译失败"
  exit 1
fi

echo "[*] Step 4: 将加密数据附加到 loader"
# 将加密数据附加到 loader 可执行文件末尾
cat combined.bin >> "$OUTPUT_LOADER"

# 写入魔数和数据大小
printf "%s" "$MAGIC" >> "$OUTPUT_LOADER"
printf "%010zu" "$PAYLOAD_LEN" >> "$OUTPUT_LOADER"

# 确保数据写入磁盘
sync

echo "[*] Step 5: 清理临时文件"
rm -f "$ENCRYPTED_BIN" "$IV_FILE" combined.bin

echo "[+] 完成！加密的 $OUTPUT_MAIN_A 已附加到 $OUTPUT_LOADER"
echo "[*] 请运行 ./$OUTPUT_LOADER 测试效果"
echo "[*] 注意: 主程序使用 AES-256-CBC 加密，密钥 '$KEY'"
echo "[*] 依赖: 确保系统已安装 OpenSSL 开发库 (sudo apt-get install libssl-dev)"
